"""
python p2_solution_bird.py
"""

import json
import numpy as np
import time
import os

class Warehouse:
    def __init__(self, warehouse_id, max_inventory, max_sales, daily_cost):
        self.warehouse_id = warehouse_id
        self.max_inventory = max_inventory
        self.max_sales = max_sales
        self.daily_cost = daily_cost

class Category:
    def __init__(self, category_id, average_inventory, average_sales):
        self.category_id = category_id
        self.average_inventory = average_inventory
        self.average_sales = average_sales

class Particle:
    def __init__(self, dimensions):
        self.position = np.zeros(dimensions, dtype=int)  # 位置
        self.velocity = np.random.randint(0, 3, size=dimensions)  # 速度
        self.best_position = self.position.copy()  # 个体最优位置
        self.best_fitness = float('inf')  # 个体最优适应度

def load_data():
    # 读取仓库信息
    with open('../fujian/fujian3/origin_data/warehouse.json') as f:
        warehouses_data = json.load(f)
    warehouses = [Warehouse(w['warehouse_id'], w['max_inventory'], w['max_sales'], w['daily_cost']) for w in warehouses_data]

    # 读取品类库存均值
    with open('../fujian/fujian3/data_from_p1/all_average_inventory.json') as f:
        inventory_data = json.load(f)
    categories = [Category(c['category_id'], c['average_inventory'], 0) for c in inventory_data]  # 这里销量暂时置为0，待后续加载

    # 读取品类销量均值
    with open('../fujian/fujian3/data_from_p1/all_average_sales.json') as f:
        sales_data = json.load(f)
    for cat in categories:
        for s in sales_data:
            if s['category_id'] == cat.category_id:
                cat.average_sales = s['average_sales']

    return warehouses, categories

def load_associations():
    with open('../fujian/fujian4/origin_data.json') as f:
        associations = json.load(f)
    association_dict = {}
    for pair in associations:
        key = (pair['category1'], pair['category2'])
        association_dict[key] = pair['association']
    return association_dict

def fitness(particle, warehouses, categories, associations):
    total_cost = 0
    inventory_utilization = 0
    sales_utilization = 0
    total_association = 0
    warehouse_inventory = {}
    warehouse_sales = {}
    
    # 计算各仓库的存货量和出货量
    for i, category in enumerate(categories):
        warehouse_id = particle.position[i]
        if warehouse_id not in warehouse_inventory:
            warehouse_inventory[warehouse_id] = 0
            warehouse_sales[warehouse_id] = 0
            
        # 计算存货量
        warehouse_inventory[warehouse_id] += category.average_inventory
        # 计算出货量
        warehouse_sales[warehouse_id] += category.average_sales

    # 计算适应度
    for warehouse in warehouses:
        if warehouse.warehouse_id in warehouse_inventory:
            inventory = warehouse_inventory[warehouse.warehouse_id]
            sales = warehouse_sales[warehouse.warehouse_id]
            total_cost += warehouse.daily_cost
            
            # 检查库存和销量限制
            if inventory < warehouse.max_inventory and sales < warehouse.max_sales:
                inventory_utilization += inventory / warehouse.max_inventory
                sales_utilization += sales / warehouse.max_sales
                
                # 计算关联度
                for i in range(len(categories)):
                    for j in range(i + 1, len(categories)):
                        if particle.position[i] == particle.position[j] == warehouse.warehouse_id:
                            key = (categories[i].category_id, categories[j].category_id)
                            total_association += associations.get(key, 0)
    
    # 适应度计算
    return total_cost * (1 + inventory_utilization) * (1 + sales_utilization) * (1 + total_association)

def update_particles(particles, global_best_position, w=0.5, c1=1.5, c2=1.5):
    for particle in particles:
        r1, r2 = np.random.random(), np.random.random()
        particle.velocity = (w * particle.velocity +
                             c1 * r1 * (particle.best_position - particle.position) +
                             c2 * r2 * (global_best_position - particle.position))
        particle.position += particle.velocity
        particle.position = np.clip(particle.position, 0, len(warehouses) - 1)  # 确保位置在合法范围

def save_best_solution(best_solution, best_fitness, time_stamp):
    solution_path = '../fujian/p2/bird/best_ever/best_solution.json'
    with open(solution_path, 'w') as f:
        json.dump([{str(i): int(best_solution[i])} for i in range(len(best_solution))], f, indent=4)

    log_path = f'../fujian/p2/bird/best_ever/best_t_{time_stamp}.log'
    with open(log_path, 'a') as f:
        f.write(f'{best_fitness}\n')

if __name__ == "__main__":
    # 加载数据
    warehouses, categories = load_data()
    associations = load_associations()

    # 初始化参数
    num_categories = len(categories)
    num_particles = 50
    particles = [Particle(num_categories) for _ in range(num_particles)]
    
    # 读取初始解
    with open('../fujian/p2/simple_solution_handwrite/ans.json') as f:
        initial_solution = json.load(f)
    for i, sol in enumerate(initial_solution):
        category_id = str(i)
        warehouse_id = sol[category_id]
        particles[i].position = np.array([warehouse_id] * num_categories)

    global_best_fitness = float('inf')
    global_best_position = None

    # 记录开始时间
    start_time = time.strftime("%Y%m%d_%H%M%S", time.localtime())
    
    try:
        while True:
            for particle in particles:
                current_fitness = fitness(particle, warehouses, categories, associations)

                # 更新个体最优
                if current_fitness < particle.best_fitness:
                    particle.best_fitness = current_fitness
                    particle.best_position = particle.position.copy()

                # 更新全局最优
                if current_fitness < global_best_fitness:
                    global_best_fitness = current_fitness
                    global_best_position = particle.position.copy()
                    save_best_solution(global_best_position, global_best_fitness, start_time)

            update_particles(particles, global_best_position)
    
    except KeyboardInterrupt:
        print("优化终止。")
